import torch
import torch.nn as nn
import torch.nn.functional as F

class CosineSimilarityLoss(nn.Module):
    def __init__(self, reduction='mean'):
        super(CosineSimilarityLoss, self).__init__()
        self.reduction = reduction

    def forward(self, input, target):
        # 计算余弦相似度
        cosine_sim = F.cosine_similarity(input, target, dim=-1)
        # 余弦相似度损失定义为 1 - 余弦相似度
        loss = 1 - cosine_sim
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:
            return loss